from keras.models import Model
from keras.applications.inception_resnet_v2 import preprocess_input
import tensorflow as tf
from data_process import load_img,img_to_array,myDataSet
from assumption import acc_top3,acc_top5
import numpy as np
import os
from classify import class_names_to_ids

# 加载模型
ws_path="/home/zhang/workspaces/cv_ws/"
# model_path="trained_models_04092348/model_05.hdf5"
model_path="trained_models_sample_04111507/model_42.hdf5"
print("Loading model :{}...".format(model_path))
customized_func={"acc_top3":acc_top3,"acc_top5":acc_top5}
model:Model=tf.keras.models.load_model(os.path.join(ws_path,model_path),customized_func)

# 创建数据加载器
imgLoader=myDataSet(data_path="/home/zhang/workspaces/cv_ws/",
                    preprocess_func=preprocess_input)

# 循环输入图片相对路径
while(True):
    img_path=input("\033[93mPath of the image (q to quit):\033[0m")
    if(img_path=="q"):
        print("Quiting...")
        break
    # 加载数据
    print("Loading image...")
    img=imgLoader.Get_Preprocess_Img([img_path])
    img=img[0]
    # 扩展维度,使其符合模型输入要求
    img:tf.Tensor=tf.expand_dims(img,axis=0)

    # 对单张图片进行预测,得到单行二维数组
    print("Predicting...")
    y:np.ndarray=model.predict(img,batch_size=1)
    # 取数组中最大值的索引作为预测类别
    y_pred_class:np.ndarray = np.argmax(y,axis=1)
    y_pred_class:np.int64=y_pred_class[0]
    for class_name in class_names_to_ids:
        if(y_pred_class == class_names_to_ids[class_name]):
            print("predict result: \033[91m{}\033[92m({} %)\033[0m"
                  .format(class_name,y[0][y_pred_class]*100))